{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac0c444c",
   "metadata": {},
   "source": [
    "# LowerTE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27718048",
   "metadata": {},
   "source": [
    "LowerTE pass 负责将 Relay 表达式转换为 TIR PrimFunc。\n",
    "\n",
    "测试用例覆盖了不同场景下的转换，包括原始函数、编译器函数和外部函数的处理。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a16dcba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbaeba3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "import tvm.testing\n",
    "import logging\n",
    "\n",
    "logging.basicConfig()\n",
    "logger = logging.getLogger(\"test_pass_lower_te\")\n",
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "# 由于 TE 编译器需要良好的重构，因此它尚未作为 '标准' pass 暴露在 relay.transform 中\n",
    "# 为了测试，直接获取该全局函数\n",
    "LowerTE = tvm._ffi.get_global_func(\"relay.tec.LowerTE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0c8f29ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform(mod):\n",
    "    \"\"\"执行转换流程，应用 LowerTE pass 到给定模块\n",
    "\n",
    "    参数:\n",
    "        mod: 输入的 TVM Relay 模块\n",
    "\n",
    "    返回:\n",
    "        tvm.ir.IRModule: 应用 LowerTE pass 后的模块\n",
    "    \"\"\"\n",
    "    logger.info(\"Starting module:\\n%s\", mod)\n",
    "    # 设置主机目标和原始目标\n",
    "    host_target = tvm.target.Target(\"llvm\")\n",
    "    prim_target = tvm.target.Target(\"llvm\", host=host_target)\n",
    "    # 创建转换上下文和编译配置\n",
    "    ctxt = tvm.transform.PassContext()\n",
    "    config = tvm.target.make_compilation_config(ctxt, prim_target)\n",
    "    # 应用设备规划转换\n",
    "    mod = tvm.relay.transform.PlanDevices(config)(mod)\n",
    "    # 应用类型推断转换\n",
    "    mod = tvm.relay.transform.InferType()(mod)\n",
    "    # 应用 LowerTE 转换，将 Relay 表达式转换为 TIR PrimFunc\n",
    "    mod = LowerTE(\"test\", config)(mod)\n",
    "    # 再次应用类型推断以确保类型正确\n",
    "    mod = tvm.relay.transform.InferType()(mod)\n",
    "    logger.info(\"After LowerTE:\\n%s\", mod)\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24bcc9e0",
   "metadata": {},
   "source": [
    "所有尝试使用结构等价测试与从 Relay 文本解析的预期 IRModule 进行比较的尝试都因设置正确的 call_lower 属性的难度而受挫，特别是设置具有正确 GlobalVar 实例的属性。因此，下面通过直接断言结构正确性来测试。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32ca0731",
   "metadata": {},
   "source": [
    "## 测试将原语 (Primitive) 函数转换为 TIR PrimFunc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "44cb1316",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:test_pass_lower_te:Starting module:\n",
      "def @main(%a: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:16 */) -> Tensor[(5, 7), float32] {\n",
      "  %0 = fn (%x: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:5:13 */, %y: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:5:17 */, Primitive=1) -> Tensor[(5, 7), float32] {\n",
      "    add(%x, %y) /* ty=Tensor[(5, 7), float32] span=from_string:5:9 */\n",
      "  } /* ty=fn (Tensor[(5, 7), float32], Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] span=from_string:7:9 */;\n",
      "  %0(%a, %a) /* ty=Tensor[(5, 7), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "INFO:test_pass_lower_te:After LowerTE:\n",
      "def @main(%a {virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a59e0640, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a5ed1a30, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))}: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:16 */, virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a59e0640, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a5ed1a30, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))) -> Tensor[(5, 7), float32] {\n",
      "  %0 = (%a, %a) /* ty=(Tensor[(5, 7), float32], Tensor[(5, 7), float32]) */;\n",
      "  call_lowered(@test_fused_add, %0, metadata={\"relay_attrs\"={__dict__={\"Primitive\"=1}}, \"all_prim_fn_vars\"=['test_fused_add']}) /* ty=Tensor[(5, 7), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "@test_fused_add = primfn(x_1: handle, y_1: handle, T_add_1: handle) -> ()\n",
      "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"test_fused_add\", \"tir.noalias\": True, \"target\": Target(id=5593a59e0640, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a5ed1a30, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}))}\n",
      "  buffers = {x: Buffer(x_2: Pointer(float32), float32, [5, 7], []),\n",
      "             y: Buffer(y_2: Pointer(float32), float32, [5, 7], []),\n",
      "             T_add: Buffer(T_add_2: Pointer(float32), float32, [5, 7], [])}\n",
      "  buffer_map = {x_1: x, y_1: y, T_add_1: T_add} {\n",
      "  for (ax0: int32, 0, 5) \"parallel\" {\n",
      "    let cse_var_1: int32 = (ax0*7)\n",
      "    T_add_3: Buffer(T_add_2, float32, [35], [])[ramp(cse_var_1, 1, 7)] = (x_3: Buffer(x_2, float32, [35], [])[ramp(cse_var_1, 1, 7)] + y_3: Buffer(y_2, float32, [35], [])[ramp(cse_var_1, 1, 7)])\n",
      "  }\n",
      "}\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>a {virtual_device<span style=\"color: #A2F; font-weight: bold\">=</span>VirtualDevice(device_type<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, virtual_device_id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>, target<span style=\"color: #A2F; font-weight: bold\">=</span>Target(id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">5593</span>a59e0640, kind<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&#39;llvm&#39;</span>, keys<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;cpu&#39;</span>}, attrs<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;mtriple&#39;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>}, host<span style=\"color: #A2F; font-weight: bold\">=</span>Target(id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">5593</span>a5ed1a30, kind<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&#39;llvm&#39;</span>, keys<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;cpu&#39;</span>}, attrs<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;mtriple&#39;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>})))}: Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">7</span>:<span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, virtual_device<span style=\"color: #A2F; font-weight: bold\">=</span>VirtualDevice(device_type<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, virtual_device_id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>, target<span style=\"color: #A2F; font-weight: bold\">=</span>Target(id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">5593</span>a59e0640, kind<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&#39;llvm&#39;</span>, keys<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;cpu&#39;</span>}, attrs<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;mtriple&#39;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>}, host<span style=\"color: #A2F; font-weight: bold\">=</span>Target(id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">5593</span>a5ed1a30, kind<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&#39;llvm&#39;</span>, keys<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;cpu&#39;</span>}, attrs<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;mtriple&#39;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>})))) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> (<span style=\"color: #A2F; font-weight: bold\">%</span>a, <span style=\"color: #A2F; font-weight: bold\">%</span>a) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>), float32], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  call_lowered(<span style=\"color: #A2F\">@test_fused_add</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, metadata<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&quot;relay_attrs&quot;</span><span style=\"color: #A2F; font-weight: bold\">=</span>{__dict__<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&quot;Primitive&quot;</span><span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>}}, <span style=\"color: #BA2121\">&quot;all_prim_fn_vars&quot;</span><span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #BA2121\">&#39;test_fused_add&#39;</span>]}) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">9</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #A2F\">@test_fused_add</span> <span style=\"color: #A2F; font-weight: bold\">=</span> primfn(x_1: handle, y_1: handle, T_add_1: handle) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> ()\n",
       "  attr <span style=\"color: #A2F; font-weight: bold\">=</span> {<span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>, <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;test_fused_add&quot;</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>, <span style=\"color: #BA2121\">&quot;target&quot;</span>: Target(id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">5593</span>a59e0640, kind<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&#39;llvm&#39;</span>, keys<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;cpu&#39;</span>}, attrs<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;mtriple&#39;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>}, host<span style=\"color: #A2F; font-weight: bold\">=</span>Target(id<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">5593</span>a5ed1a30, kind<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&#39;llvm&#39;</span>, keys<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;cpu&#39;</span>}, attrs<span style=\"color: #A2F; font-weight: bold\">=</span>{<span style=\"color: #BA2121\">&#39;mtriple&#39;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>}))}\n",
       "  buffers <span style=\"color: #A2F; font-weight: bold\">=</span> {x: Buffer(x_2: Pointer(float32), float32, [<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>], []),\n",
       "             y: Buffer(y_2: Pointer(float32), float32, [<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>], []),\n",
       "             T_add: Buffer(T_add_2: Pointer(float32), float32, [<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">7</span>], [])}\n",
       "  buffer_map <span style=\"color: #A2F; font-weight: bold\">=</span> {x_1: x, y_1: y, T_add_1: T_add} {\n",
       "  <span style=\"color: #008000; font-weight: bold\">for</span> (ax0: int32, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>) <span style=\"color: #BA2121\">&quot;parallel&quot;</span> {\n",
       "    let cse_var_1: int32 <span style=\"color: #A2F; font-weight: bold\">=</span> (ax0<span style=\"color: #A2F; font-weight: bold\">*</span><span style=\"color: #008000\">7</span>)\n",
       "    T_add_3: Buffer(T_add_2, float32, [<span style=\"color: #008000\">35</span>], [])[ramp(cse_var_1, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>)] <span style=\"color: #A2F; font-weight: bold\">=</span> (x_3: Buffer(x_2, float32, [<span style=\"color: #008000\">35</span>], [])[ramp(cse_var_1, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>)] <span style=\"color: #A2F; font-weight: bold\">+</span> y_3: Buffer(y_2, float32, [<span style=\"color: #008000\">35</span>], [])[ramp(cse_var_1, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>)])\n",
       "  }\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 创建包含原始函数的输入模块\n",
    "input_mod = tvm.relay.parse(\n",
    "    \"\"\"\n",
    "    #[version = \"0.0.5\"]\n",
    "    def @main(%a: Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] {\n",
    "        %0 = fn(%x : Tensor[(5, 7), float32], %y : Tensor[(5, 7), float32], Primitive=1) -> Tensor[(5, 7), float32] {\n",
    "        add(%x, %y)\n",
    "        };\n",
    "        %0(%a, %a)\n",
    "    }\n",
    "    \"\"\",\n",
    "    \"from_string\",\n",
    "    None,\n",
    "    None,\n",
    ")\n",
    "\n",
    "# 应用转换\n",
    "actual_mod = transform(input_mod)\n",
    "actual_mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e227e01c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证主函数结构\n",
    "main = actual_mod[\"main\"]\n",
    "call = main.body\n",
    "assert call.op.name == \"call_lowered\"\n",
    "assert len(call.args) == 2\n",
    "assert call.args[0].name_hint == \"test_fused_add\"\n",
    "assert len(call.args[1].fields) == 2\n",
    "assert call.args[1].fields[0].name_hint == \"a\"\n",
    "assert call.args[1].fields[1].name_hint == \"a\"\n",
    "assert call.attrs.metadata[\"relay_attrs\"].Primitive == 1\n",
    "assert len(call.attrs.metadata[\"all_prim_fn_vars\"]) == 1\n",
    "assert call.attrs.metadata[\"all_prim_fn_vars\"][0].name_hint == \"test_fused_add\"\n",
    "# 验证生成的 PrimFunc\n",
    "test_fused_add = actual_mod[\"test_fused_add\"]\n",
    "assert isinstance(test_fused_add, tvm.tir.PrimFunc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef7ec5c5",
   "metadata": {},
   "source": [
    "## 测试将带有编译器标记的函数转换为外部函数\n",
    "\n",
    "验证标记为 Compiler=\"test_pass_lower_te\" 的函数能否被正确处理并转换为外部函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5d340112",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:test_pass_lower_te:Starting module:\n",
      "def @main(%a: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:16 */) -> Tensor[(5, 7), float32] {\n",
      "  %0 = fn (%x: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:5:13 */, %y: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:5:17 */, Primitive=1, Compiler=\"test_pass_lower_te\", global_symbol=\"test_add\") -> Tensor[(5, 7), float32] {\n",
      "    add(%x, %y) /* ty=Tensor[(5, 7), float32] span=from_string:5:9 */\n",
      "  } /* ty=fn (Tensor[(5, 7), float32], Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] span=from_string:7:9 */;\n",
      "  %0(%a, %a) /* ty=Tensor[(5, 7), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "INFO:test_pass_lower_te:After LowerTE:\n",
      "def @main(%a {virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a5a0b1f0, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a5bd4bf0, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))}: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:16 */, virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a5a0b1f0, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a5bd4bf0, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))) -> Tensor[(5, 7), float32] {\n",
      "  %0 = (%a, %a) /* ty=(Tensor[(5, 7), float32], Tensor[(5, 7), float32]) */;\n",
      "  call_lowered(@test_add, %0, metadata={\"relay_attrs\"={__dict__={\"Primitive\"=1, \"Compiler\"=\"test_pass_lower_te\", \"global_symbol\"=\"test_add\"}}, \"all_prim_fn_vars\"=[]}) /* ty=Tensor[(5, 7), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "def @test_add(%x: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:5:13 */, %y: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:5:17 */, Extern=1) -> Tensor[(5, 7), float32] {\n",
      "  add(%x, %y) /* ty=Tensor[(5, 7), float32] span=from_string:5:9 */\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 注册测试用的外部函数\n",
    "@tvm._ffi.register_func(\"relay.ext.test_pass_lower_te\")\n",
    "def relay_ext_test_pass_lower_te(func):\n",
    "    return None\n",
    "\n",
    "# 创建包含编译器函数的输入模块\n",
    "input_mod = tvm.relay.parse(\n",
    "    \"\"\"\n",
    "    #[version = \"0.0.5\"]\n",
    "    def @main(%a: Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] {\n",
    "        %0 = fn(%x : Tensor[(5, 7), float32], %y : Tensor[(5, 7), float32], Primitive=1, Compiler=\"test_pass_lower_te\", global_symbol=\"test_add\") -> Tensor[(5, 7), float32] {\n",
    "        add(%x, %y)\n",
    "        };\n",
    "        %0(%a, %a)\n",
    "    }\n",
    "    \"\"\",\n",
    "    \"from_string\",\n",
    "    None,\n",
    "    None,\n",
    ")\n",
    "\n",
    "# 应用转换\n",
    "actual_mod = transform(input_mod)\n",
    "\n",
    "# 预期结果:\n",
    "#   def @main(%a : Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] {\n",
    "#     %0 = (%a, %a)\n",
    "#     call_lowered(@test_add , %0, metadata={relay_attrs={Primitive=1, Compiler=\"test_pass_lower_te\", global_symbol=\"test_add\"}}, all_prim_fn_vars=[]})\n",
    "#   }\n",
    "#   def @test_add(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32], Extern=1) -> Tensor[(5, 7), float32] {\n",
    "#     add(%x, %y)\n",
    "#   }\n",
    "\n",
    "# 验证主函数结构\n",
    "main = actual_mod[\"main\"]\n",
    "call = main.body\n",
    "assert call.op.name == \"call_lowered\"\n",
    "assert len(call.args) == 2\n",
    "assert call.args[0].name_hint == \"test_add\"\n",
    "assert len(call.args[1].fields) == 2\n",
    "assert call.args[1].fields[0].name_hint == \"a\"\n",
    "assert call.args[1].fields[1].name_hint == \"a\"\n",
    "assert call.attrs.metadata[\"relay_attrs\"].Primitive == 1\n",
    "assert call.attrs.metadata[\"relay_attrs\"].Compiler == \"test_pass_lower_te\"\n",
    "assert call.attrs.metadata[\"relay_attrs\"].global_symbol == \"test_add\"\n",
    "assert len(call.attrs.metadata[\"all_prim_fn_vars\"]) == 0\n",
    "\n",
    "# 验证生成的外部函数\n",
    "test_add = actual_mod[\"test_add\"]\n",
    "assert isinstance(test_add, tvm.relay.Function)\n",
    "assert test_add.attrs[\"Extern\"] == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bd21f4b",
   "metadata": {},
   "source": [
    "## 测试处理已标记为外部的函数\n",
    "验证已经标记为 Extern=1 的全局函数能否被正确处理："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f0d1a3d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:test_pass_lower_te:Starting module:\n",
      "def @main(%a: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:4:22 */) -> Tensor[(5, 7), float32] {\n",
      "  @my_add(%a, %a) /* ty=Tensor[(5, 7), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "def @my_add(%x: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:13 */, %y: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:17 */, Extern=1) -> Tensor[(5, 7), float32] {\n",
      "  add(%x, %y) /* ty=Tensor[(5, 7), float32] span=from_string:7:9 */\n",
      "}\n",
      "\n",
      "INFO:test_pass_lower_te:After LowerTE:\n",
      "def @main(%a {virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a574e8c0, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a5725150, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))}: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:4:22 */, virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a574e8c0, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a5725150, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))) -> Tensor[(5, 7), float32] {\n",
      "  %0 = (%a, %a) /* ty=(Tensor[(5, 7), float32], Tensor[(5, 7), float32]) */;\n",
      "  call_lowered(@my_add, %0, metadata={\"relay_attrs\"={__dict__={\"Extern\"=1}}, \"all_prim_fn_vars\"=[]}) /* ty=Tensor[(5, 7), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "def @my_add(%x: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:13 */, %y: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:17 */, Extern=1) -> Tensor[(5, 7), float32] {\n",
      "  add(%x, %y) /* ty=Tensor[(5, 7), float32] span=from_string:7:9 */\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 创建包含外部函数的输入模块\n",
    "input_mod = tvm.relay.parse(\n",
    "    \"\"\"\n",
    "    #[version = \"0.0.5\"]\n",
    "    def @main(%a: Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] {\n",
    "        @my_add(%a, %a)\n",
    "    }\n",
    "    def @my_add(%x : Tensor[(5, 7), float32], %y : Tensor[(5, 7), float32], Extern=1) -> Tensor[(5, 7), float32] {\n",
    "        add(%x, %y)\n",
    "    }\n",
    "    \"\"\",\n",
    "    \"from_string\",\n",
    "    None,\n",
    "    None,\n",
    ")\n",
    "\n",
    "# 应用转换\n",
    "actual_mod = transform(input_mod)\n",
    "\n",
    "# 预期结果:\n",
    "#   def @main(%a: Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] {\n",
    "#     %0 = (%a, %a);\n",
    "#     call_lowered(@my_add, %0, metadata={relay_attrs={Extern=1}}, all_prim_fn_vars=[]})\n",
    "#   }\n",
    "#   def @my_add(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32], Extern=1) -> Tensor[(5, 7), float32] {\n",
    "#     add(%x, %y)\n",
    "#   }\n",
    "\n",
    "# 验证主函数结构\n",
    "main = actual_mod[\"main\"]\n",
    "call = main.body\n",
    "assert call.op.name == \"call_lowered\"\n",
    "assert len(call.args) == 2\n",
    "assert call.args[0].name_hint == \"my_add\"\n",
    "assert len(call.args[1].fields) == 2\n",
    "assert call.args[1].fields[0].name_hint == \"a\"\n",
    "assert call.args[1].fields[1].name_hint == \"a\"\n",
    "assert call.attrs.metadata[\"relay_attrs\"].Extern == 1\n",
    "assert len(call.attrs.metadata[\"all_prim_fn_vars\"]) == 0\n",
    "\n",
    "# 验证外部函数保持不变\n",
    "test_add = actual_mod[\"my_add\"]\n",
    "assert isinstance(test_add, tvm.relay.Function)\n",
    "assert test_add.attrs[\"Extern\"] == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b9591d1",
   "metadata": {},
   "source": [
    "## 测试处理具有动态形状的外部函数\n",
    "\n",
    "验证带有动态形状（`?` 维度）的外部函数能否被正确处理，并生成相应的形状函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a3e63556",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:test_pass_lower_te:Starting module:\n",
      "def @main(%a: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:4:22 */) -> Tensor[(?, ?), float32] {\n",
      "  @my_dyn(%a, %a) /* ty=Tensor[(?, ?), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "def @my_dyn(%x: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:13 */, %y: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:17 */, Extern=1) -> Tensor[(?, ?), float32] {\n",
      "  add(%x, %y) /* ty=Tensor[(?, ?), float32] span=from_string:7:9 */\n",
      "}\n",
      "\n",
      "INFO:test_pass_lower_te:After LowerTE:\n",
      "def @main(%a {virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a598f300, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a595af10, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))}: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:4:22 */, virtual_device=VirtualDevice(device_type=1, virtual_device_id=0, target=Target(id=5593a598f300, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"}, host=Target(id=5593a595af10, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})))) -> Tensor[(?, ?), float32] {\n",
      "  %0 = (%a, %a) /* ty=(Tensor[(5, 7), float32], Tensor[(5, 7), float32]) */;\n",
      "  call_lowered(@my_dyn, %0, metadata={\"prim_shape_fn_num_outputs\"=1, \"prim_shape_fn_num_inputs\"=2, \"all_prim_fn_vars\"=[], \"relay_attrs\"={__dict__={\"Extern\"=1}}, \"prim_shape_fn_var\"='test_shape_func_add', \"prim_shape_fn_states\"=[2, 2], \"all_prim_shape_fn_vars\"=['test_shape_func_add']}) /* ty=Tensor[(?, ?), float32] span=from_string:4:9 */\n",
      "}\n",
      "\n",
      "def @my_dyn(%x: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:13 */, %y: Tensor[(5, 7), float32] /* ty=Tensor[(5, 7), float32] span=from_string:7:17 */, Extern=1) -> Tensor[(?, ?), float32] {\n",
      "  add(%x, %y) /* ty=Tensor[(?, ?), float32] span=from_string:7:9 */\n",
      "}\n",
      "\n",
      "@test_shape_func_add = primfn(shape_x_1: handle, shape_y_1: handle, v_broadcast_shape_func: handle) -> ()\n",
      "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"test_shape_func_add\", \"tir.noalias\": True, \"target\": Target(id=5593a595af10, kind='llvm', keys={'cpu'}, attrs={'mtriple': \"x86_64-unknown-linux-gnu\"})}\n",
      "  buffers = {shape_x: Buffer(shape_x_2: Pointer(int64), int64, [2], []),\n",
      "             shape_y: Buffer(shape_y_2: Pointer(int64), int64, [2], []),\n",
      "             buf__broadcast_shape_func: Buffer(v_broadcast_shape_func_1: Pointer(int64), int64, [2], [])}\n",
      "  buffer_map = {shape_x_1: shape_x, shape_y_1: shape_y, v_broadcast_shape_func: buf__broadcast_shape_func} {\n",
      "  attr [0] \"extern_scope\" = 0 {\n",
      "    if (shape_x[1] == shape_y[1]) {\n",
      "      buf__broadcast_shape_func[1] = shape_x[1]\n",
      "    } else {\n",
      "      if (shape_x[1] == 1i64) {\n",
      "        buf__broadcast_shape_func[1] = shape_y[1]\n",
      "      } else {\n",
      "        assert((shape_y[1] == 1i64), \"Incompatible broadcast type shape_x[1] and shape_y[1]\")\n",
      "        0\n",
      "        buf__broadcast_shape_func[1] = shape_x[1]\n",
      "      }\n",
      "    }\n",
      "    if (shape_x[0] == shape_y[0]) {\n",
      "      buf__broadcast_shape_func[0] = shape_x[0]\n",
      "    } else {\n",
      "      if (shape_x[0] == 1i64) {\n",
      "        buf__broadcast_shape_func[0] = shape_y[0]\n",
      "      } else {\n",
      "        assert((shape_y[0] == 1i64), \"Incompatible broadcast type shape_x[0] and shape_y[0]\")\n",
      "        0\n",
      "        buf__broadcast_shape_func[0] = shape_x[0]\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 创建包含动态形状外部函数的输入模块\n",
    "input_mod = tvm.relay.parse(\n",
    "    \"\"\"\n",
    "    #[version = \"0.0.5\"]\n",
    "    def @main(%a: Tensor[(5, 7), float32]) -> Tensor[(?, ?), float32] {\n",
    "        @my_dyn(%a, %a)\n",
    "    }\n",
    "    def @my_dyn(%x : Tensor[(5, 7), float32], %y : Tensor[(5, 7), float32], Extern=1) -> Tensor[(?, ?), float32] {\n",
    "        add(%x, %y)\n",
    "    }\n",
    "    \"\"\",\n",
    "    \"from_string\",\n",
    "    None,\n",
    "    None,\n",
    ")\n",
    "\n",
    "# 应用转换\n",
    "actual_mod = transform(input_mod)\n",
    "\n",
    "# 预期结果:\n",
    "# def @main(%a: Tensor[(5, 7), float32]) -> Tensor[(?, ?), float32] {\n",
    "#   %0 = (%a, %a);\n",
    "#   call_lowered(@my_dyn, %0, metadata={prim_shape_fn_var='test_shape_func_add', relay_attrs={Extern=1}, prim_shape_fn_states=[2, 2], prim_shape_fn_num_inputs=2, all_prim_shape_fn_vars=['shape_func_add'], prim_shape_fn_num_outputs=1, all_prim_fn_vars=[]})\n",
    "# }\n",
    "# def @my_dyn(%x: Tensor[(5, 7), float32] , %y: Tensor[(5, 7), float32] , Extern=1) -> Tensor[(?, ?), float32] {\n",
    "#   add(%x, %y)\n",
    "# }\n",
    "# def @test_shape_func_add = <shape PrimFunc>\n",
    "\n",
    "# 验证主函数结构和动态形状相关的元数据\n",
    "main = actual_mod[\"main\"]\n",
    "call = main.body\n",
    "assert call.op.name == \"call_lowered\"\n",
    "assert len(call.args) == 2\n",
    "assert call.args[0].name_hint == \"my_dyn\"\n",
    "assert len(call.args[1].fields) == 2\n",
    "assert call.args[1].fields[0].name_hint == \"a\"\n",
    "assert call.args[1].fields[1].name_hint == \"a\"\n",
    "assert call.attrs.metadata[\"prim_shape_fn_var\"].name_hint == \"test_shape_func_add\"\n",
    "assert call.attrs.metadata[\"relay_attrs\"].Extern == 1\n",
    "assert len(call.attrs.metadata[\"prim_shape_fn_states\"]) == 2\n",
    "assert call.attrs.metadata[\"prim_shape_fn_states\"][0] == 2\n",
    "assert call.attrs.metadata[\"prim_shape_fn_states\"][1] == 2\n",
    "assert call.attrs.metadata[\"prim_shape_fn_num_inputs\"] == 2\n",
    "assert len(call.attrs.metadata[\"all_prim_shape_fn_vars\"]) == 1\n",
    "assert call.attrs.metadata[\"all_prim_shape_fn_vars\"][0].name_hint == \"test_shape_func_add\"\n",
    "assert call.attrs.metadata[\"prim_shape_fn_num_outputs\"] == 1\n",
    "assert len(call.attrs.metadata[\"all_prim_fn_vars\"]) == 0\n",
    "\n",
    "# 验证外部函数保持不变\n",
    "my_dyn = actual_mod[\"my_dyn\"]\n",
    "assert isinstance(my_dyn, tvm.relay.Function)\n",
    "assert my_dyn.attrs[\"Extern\"] == 1\n",
    "\n",
    "# 验证生成的形状函数\n",
    "shape_func_add = actual_mod[\"test_shape_func_add\"]\n",
    "assert isinstance(shape_func_add, tvm.tir.PrimFunc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4daadb7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
